import anndata as ad
import scanpy as sc
import sctm
import squidpy as sq
from sklearn.decomposition import NMF, PCA

from pyfeasc.utils.mca import run_mca


def scamp(adata: ad.AnnData,
          n_topics: int = 16,
          batch: str = "batch",
          method: str = 'mca',
          n_comps: int = 50,
          n_comps_1: int = 40,
          flag: bool = False,
          device: str = "cpu"):
	"""
	功能：降维、主题模型(使用原始 count 矩阵)
	:param adata: anndata 对象
	:param method: 降维方法，可选 'pca', 'nmf','mca','mca_pca'
	:param n_comps: 第一次降维 （pca, nmf, mca) 后的维数
	:param n_comps_1: 第二次降维后的维数 （如果 method='mca_pca'）
	:param n_topics: stamp 后的主题数
	:param batch: 如果有分批次数据，则指定 batch 列名
	:param flag: 是否分批次降维
	:param device: 选择运行设备，'cpu' 或 'cuda:0'
	:return: topic_prop：每个细胞的主题分布，(n_cells, n_topics)，beta：每个特征的主题分布，(n_hv_genes, n_topics)
	"""
	
	adata.layers['raw'] = adata.X.copy()
	
	# 归一化（不能使用 sc.pp.scale）
	sc.pp.normalize_total(adata, target_sum=1e4)
	sc.pp.log1p(adata)
	
	if flag:
		batch_list = list(adata.obs[batch].unique())
		adata_batch_list = [adata[adata.obs[batch] == b] for b in batch_list]
		
		if method == 'pca':
			for adata_batch in adata_batch_list:
				sc.pp.pca(adata_batch, n_comps=n_comps)
			adata = adata_batch_list[0].concatenate(adata_batch_list[1:], batch_key="batch_list", index_unique=None)
			adata.obsm['spatial'] = adata.obsm['X_pca']
		elif method == 'nmf':
			for adata_batch in adata_batch_list:
				nmf = NMF(n_components=n_comps)
				adata_batch.obsm['X_nmf'] = nmf.fit_transform(adata_batch.X)
			adata = adata_batch_list[0].concatenate(adata_batch_list[1:], batch_key="batch_list", index_unique=None)
			adata.obsm['spatial'] = adata.obsm['X_nmf']
		elif method == 'mca':
			for adata_batch in adata_batch_list:
				out = run_mca(adata_batch, nmcs=n_comps)
				adata_batch.obsm['X_mca'] = out[0]
			adata = adata_batch_list[0].concatenate(adata_batch_list[1:], batch_key="batch_list", index_unique=None)
			adata.obsm['spatial'] = adata.obsm['X_mca']
		elif method == 'mca_pca':
			for adata_batch in adata_batch_list:
				out = run_mca(adata_batch, nmcs=n_comps)
				adata_batch.obsm['X_mca'] = out[0]
			for adata_batch in adata_batch_list:
				pca = PCA(n_components=n_comps_1)
				adata_batch.obsm['X_pca'] = pca.fit_transform(adata_batch.obsm['X_mca'])
			adata = adata_batch_list[0].concatenate(adata_batch_list[1:], batch_key="batch_list", index_unique=None)
			adata.obsm['spatial'] = adata.obsm['X_pca']
		else:
			raise ValueError("Invalid method. Choose from 'pca', 'nmf', 'mca', or'mca_pca'.")
	else:
		if method == 'nmf':
			nmf = NMF(n_components=n_comps)
			adata.obsm['X_nmf'] = nmf.fit_transform(adata.X)
			adata.obsm['spatial'] = adata.obsm['X_nmf']
		elif method == 'pca':
			sc.tl.pca(adata, n_comps=n_comps)
			adata.obsm['spatial'] = adata.obsm['X_pca']
		elif method == 'mca':
			out = run_mca(adata, nmcs=n_comps)
			adata.obsm['spatial'] = out[0]
		elif method == 'mca_pca':
			out = run_mca(adata, nmcs=n_comps)
			pca = PCA(n_components=n_comps_1)
			adata.obsm['X_pca'] = pca.fit_transform(out[0])
			adata.obsm['spatial'] = adata.obsm['X_pca']
		else:
			raise ValueError("Invalid method. Choose from 'pca', 'mca', 'nmf', or 'mca_pca'.")
	
	sq.gr.spatial_neighbors(adata)
	adata.X = adata.layers['raw']
	
	if batch:
		model = sctm.stamp.STAMP(
				adata,
				n_topics=n_topics,
				categorical_covariate_keys=[batch],
				gene_likelihood="nb",
				dropout=0.1,
		)
	else:
		model = sctm.stamp.STAMP(
				adata,
				n_topics=n_topics,
		)
	model.train(batch_size=4096, device=device)
	
	topic_prop = model.get_cell_by_topic()  # (n_cells, n_topics)
	beta = model.get_feature_by_topic()  # (n_hv_genes, n_topics)
	
	return topic_prop, beta
